package com.ihai.spring.boot.sample.es;

import com.alibaba.fastjson.JSON;
import com.ihai.spring.boot.sample.es.config.EsConfiguration;
import org.apache.http.HttpHost;
import org.apache.http.impl.nio.conn.PoolingNHttpClientConnectionManager;
import org.apache.http.impl.nio.reactor.DefaultConnectingIOReactor;
import org.apache.http.impl.nio.reactor.IOReactorConfig;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.DocWriteRequest;
import org.elasticsearch.action.admin.indices.create.CreateIndexRequest;
import org.elasticsearch.action.admin.indices.create.CreateIndexResponse;
import org.elasticsearch.action.admin.indices.get.GetIndexRequest;
import org.elasticsearch.action.admin.indices.mapping.put.PutMappingRequest;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.support.master.AcknowledgedResponse;
import org.elasticsearch.client.*;
import org.elasticsearch.common.io.stream.BytesStreamOutput;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentFactory;
import org.elasticsearch.common.xcontent.XContentType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.BufferedOutputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Calendar;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

@Component
public class EsUtil {
    private static final Logger logger = LoggerFactory.getLogger(EsUtil.class);
    private static final String TYPE = "_doc";
    @Autowired
    private EsConfiguration esConfiguration;

    private Object lock = new Object();
    private RestHighLevelClient client;
    private AtomicInteger i = new AtomicInteger(0);

    private RestHighLevelClient getClient() {
        if(client == null){
            synchronized (lock) {
                if(client == null) {
                    List<HttpHost> httpHostList = new ArrayList<>();
                    for (String uri : esConfiguration.getUris()) {
                        httpHostList.add(HttpHost.create(uri));
                    }
                    try {
                        IOReactorConfig ioReactorConfig = IOReactorConfig.custom().setIoThreadCount(100).setConnectTimeout(10).setSoKeepAlive(true).build();
                        PoolingNHttpClientConnectionManager connManager = new PoolingNHttpClientConnectionManager(new DefaultConnectingIOReactor(ioReactorConfig));
                        connManager.setMaxTotal(10000);
                        connManager.setDefaultMaxPerRoute(10000);
                        RestClientBuilder restClientBuilder = RestClient.builder(httpHostList.toArray(new HttpHost[0]))
                                .setHttpClientConfigCallback(httpClientBuilder -> {
                                    return httpClientBuilder.setConnectionManager(connManager);
                                })
                                .setMaxRetryTimeoutMillis(5 * 60 * 1000);
                        client = new RestHighLevelClient(restClientBuilder);
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                }
            }
        }
        return client;
    }

    public void createIndex(String indexPrefix, String index) {
        try {
            CreateIndexRequest createIndexRequest = new CreateIndexRequest(index);
            createIndexRequest.settings(Settings.builder().put("number_of_shards", 1).put("number_of_replicas",0).put("refresh_interval","-1"));
            XContentBuilder builder = XContentFactory.jsonBuilder();
            builder.startObject();
            builder.startObject("properties");
            for(String property : esConfiguration.getMapping().get(indexPrefix).keySet()){
                builder.startObject(property);
                Map<String,String> typeMap = esConfiguration.getMapping().get(indexPrefix).get(property);
                for(String key : typeMap.keySet()) {
                    builder.field(key, typeMap.get(key));
                }
                builder.endObject();
            }
            builder.endObject();
            builder.endObject();
            createIndexRequest.mapping(TYPE, builder);
            getClient().indices().createAsync(createIndexRequest, RequestOptions.DEFAULT, new ActionListener<CreateIndexResponse>() {
                @Override
                public void onResponse(CreateIndexResponse createIndexResponse) {
                    ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
                    BytesStreamOutput bytesStreamOutput = new BytesStreamOutput();
                    try {
                        createIndexResponse.writeTo(bytesStreamOutput);
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                    logger.info(bytesStreamOutput.toString());
                }

                @Override
                public void onFailure(Exception e) {
                    logger.error("create index error.", e);
                }
            });
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public IndexRequest createIndexRequest(String index, String contentType, String opType, Object source) {
        IndexRequest indexRequest = Requests.indexRequest(index)
                .type(TYPE)
                .opType(opType);
        if(source != null) {
            indexRequest.source(JSON.toJSON(source).toString(), XContentType.valueOf(contentType));
        }
        return indexRequest;
    }

    public IndexRequest createIndexRequest(String index, Object source) {
        return createIndexRequest(index, XContentType.JSON.name(), DocWriteRequest.OpType.INDEX.name(),source);
    }

    public boolean existsIndex(String index) {
        try {
            GetIndexRequest getIndexRequest = new GetIndexRequest();
            getIndexRequest.indices(index);
            return getClient().indices().exists(getIndexRequest, RequestOptions.DEFAULT);
        }catch (IOException e) {
            e.printStackTrace();
            return false;
        }
    }

    public void bulk(String indexPrefix, List<?> sources) {
        if(sources == null || sources.size() == 0) {
            return;
        }
        Calendar cal = Calendar.getInstance();
        String index = indexPrefix + cal.get(Calendar.YEAR) + "-" + (cal.get(Calendar.MONTH) + 1);
        if(!existsIndex(index)) {
            createIndex(indexPrefix, index);
        }
        BulkRequest bulkRequest = Requests.bulkRequest();
        for(Object source : sources) {
            bulkRequest.add(createIndexRequest(index, source));
        }
        try {
            long startTime = System.currentTimeMillis();
            getClient().bulk(bulkRequest, RequestOptions.DEFAULT);
            logger.info("ElasticSearch bulk " + sources.size() + " records costs " + (System.currentTimeMillis() - startTime) + "ms.");
        }catch (IOException e){
            logger.error("ElasticSearch bulk failed。", e);
        }
    }

    public void bulkAsync(String indexPrefix, List<?> sources) {
        int j = i.incrementAndGet();
        logger.info("-----------"+j+"--");
        if(sources == null || sources.size() == 0) {
            return;
        }
        Calendar cal = Calendar.getInstance();
        String index = indexPrefix + cal.get(Calendar.YEAR) + "-" + (cal.get(Calendar.MONTH) + 1);
        if(!existsIndex(index)) {
            createIndex(indexPrefix, index);
        }
        long startTime = System.currentTimeMillis();
        BulkRequest bulkRequest = Requests.bulkRequest();
        for(Object source : sources) {
            bulkRequest.add(createIndexRequest(index, source));
        }
        logger.info("-----------"+j+"--");
        getClient().bulkAsync(bulkRequest, RequestOptions.DEFAULT,new ActionListener<BulkResponse>() {

            @Override
            public void onResponse(BulkResponse response) {
                logger.info(j+"ElasticSearch bulk "+sources.size()+" records costs " + (System.currentTimeMillis() - startTime) +"ms.");
            }

            @Override
            public void onFailure(Exception e) {
                logger.error("ElasticSearch bulk failed。", e);
            }
        });
        logger.info("-----------"+j+"--");
    }
}
